#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
检查 RL Parquet 中“正常眼底”样本数量
====================================

【功能说明】
- 读入一个 RL 阶段的 Parquet 数据文件。
- 根据全局 BLACK_LIST 判定样本是否为“正常”。
- 输出：总数、正常数目、比例。

【用法示例】
python fundus_reasoner/data_preprocess/1_1_check_normal.py \
  --input_parquet ./experiments/dataset/RL_parquet/med_pub_rl_test.parquet
"""

import argparse
import os
import re
from typing import List

import pandas as pd

# ===== 全局常量 =====
BLACK_LIST = ["正常眼底"]
SEP_PATTERN = r"[,\s，、]+"  # 标签切分分隔符


def parse_args():
    parser = argparse.ArgumentParser(description="检查 Parquet 中正常样本比例")
    parser.add_argument(
        "--input_parquet",
        type=str,
        required=True,
        help="输入的 Parquet 文件路径",
    )
    return parser.parse_args()


def extract_labels(ground_truth: str) -> List[str]:
    """从 ground_truth 中提取标签"""
    if not isinstance(ground_truth, str):
        return []
    if "诊断为" in ground_truth:
        tail = ground_truth.split("诊断为", 1)[-1]
    else:
        tail = ground_truth
    parts = re.split(SEP_PATTERN, tail.strip())
    return [p.strip() for p in parts if p.strip()]


def is_normal(row) -> bool:
    """判定是否正常样本（answer 或 ground_truth 中包含 BLACK_LIST）"""
    rm = row.get("reward_model", None)
    gt = ""
    if isinstance(rm, dict):
        gt = rm.get("ground_truth", "")
    labels = extract_labels(gt)
    # 判断是否包含黑名单关键词
    return any(lbl in BLACK_LIST for lbl in labels)


def main():
    args = parse_args()
    in_path = args.input_parquet
    assert os.path.exists(in_path), f"输入文件不存在：{in_path}"

    df = pd.read_parquet(in_path)
    total = len(df)

    normal_mask = df.apply(is_normal, axis=1)
    n_normal = int(normal_mask.sum())
    percent = (n_normal / total * 100) if total > 0 else 0.0

    print(
        f"共检测到数据 {total} 条，其中正常样本 {n_normal} 条，占比 {percent:.2f}%"
    )


if __name__ == "__main__":
    main()
